{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import pprint\n",
    "pp = pprint.PrettyPrinter(indent=4)\n",
    "from keras import backend as K\n",
    "from keras.models import Sequential, Model, load_model\n",
    "from keras.layers.embeddings import Embedding\n",
    "from keras.layers import Input, Activation, Dense, Lambda, Permute, Dropout, add, dot, concatenate, multiply\n",
    "from keras.layers import LSTM\n",
    "from keras.utils.data_utils import get_file\n",
    "from keras.preprocessing.sequence import pad_sequences\n",
    "from functools import reduce\n",
    "import tarfile\n",
    "import numpy as np\n",
    "import re\n",
    "import tensorflow as tf\n",
    "from nltk.tokenize import word_tokenize\n",
    "from keras.callbacks import ModelCheckpoint\n",
    "import datetime\n",
    "from itertools import chain\n",
    "from process_data import load_entities, save_pickle, load_pickle, load_kv_pairs, lower_list, vectorize, vectorize_kv, find_ngrams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load pickle/mov_task1_qa_pipe_train.pickle\n",
      "load pickle/mov_task1_qa_pipe_test.pickle\n",
      "load pickle/mov_kv_pairs.pickle\n",
      "load pickle/mov_train_k.pickle\n",
      "load pickle/mov_train_v.pickle\n",
      "load pickle/mov_test_k.pickle\n",
      "load pickle/mov_test_v.pickle\n",
      "load pickle/mov_entities.pickle\n",
      "load pickle/mov_vocab.pickle\n",
      "load pickle/mov_stopwords.pickle\n",
      "load pickle/mov_w2i.pickle\n",
      "load pickle/mov_i2w.pickle\n",
      "-\n",
      "Vocab size: 74609 unique words\n",
      "Query max length: 16 words\n",
      "Number of training data: 96185\n",
      "Number of test data: 9952\n",
      "-\n",
      "Here's what a \"data\" tuple looks like (input, query, answer):\n",
      "([], ['what', 'movies', 'are', 'about', 'ginger rogers'], ['the barkleys of broadway', 'kitty foyle', 'top hat'])\n"
     ]
    }
   ],
   "source": [
    "is_babi = False\n",
    "if is_babi:\n",
    "    train_data = load_task('./data/tasks_1-20_v1-2/en/qa5_three-arg-relations_train.txt', is_babi)\n",
    "    test_data = load_task('./data/tasks_1-20_v1-2/en/qa5_three-arg-relations_test.txt', is_babi)\n",
    "else:\n",
    "    # mem_maxlen         = 100 # 1つのエピソードに関連しているKVの数に対する制限\n",
    "    train_data         = load_pickle('pickle/mov_task1_qa_pipe_train.pickle')\n",
    "    test_data          = load_pickle('pickle/mov_task1_qa_pipe_test.pickle')\n",
    "    kv_pairs           = load_pickle('pickle/mov_kv_pairs.pickle')\n",
    "    train_k            = np.array(load_pickle('pickle/mov_train_k.pickle'))\n",
    "    train_v            = np.array(load_pickle('pickle/mov_train_v.pickle'))\n",
    "    test_k             = np.array(load_pickle('pickle/mov_test_k.pickle'))\n",
    "    test_v             = np.array(load_pickle('pickle/mov_test_v.pickle'))\n",
    "    entities           = load_pickle('pickle/mov_entities.pickle')\n",
    "    entity_size        = len(entities)\n",
    "\n",
    "# vocab = set(entities + ['directed_by', 'written_by', 'starred_actors', 'release_year', 'has_genre', 'has_tags', 'has_plot'] \n",
    "#                      + ['!directed_by', '!written_by', '!starred_actors', '!release_year', '!has_genre', '!has_tags', '!has_plot'] )\n",
    "# for _, q, answer in train_data + test_data:\n",
    "#     vocab |= set(q + answer)\n",
    "# vocab = sorted(vocab)\n",
    "# save_pickle(vocab, 'pickle/mov_vocab.pickle')\n",
    "vocab = load_pickle('pickle/mov_vocab.pickle')\n",
    "stopwords = load_pickle('pickle/mov_stopwords.pickle')\n",
    "\n",
    "# w2i = dict((c, i) for i, c in enumerate(vocab))\n",
    "# i2w = dict((i, c) for i, c in enumerate(vocab))\n",
    "# save_pickle(w2i, 'mov_w2i.pickle')\n",
    "# save_pickle(i2w, 'mov_i2w.pickle')\n",
    "w2i = load_pickle('pickle/mov_w2i.pickle')\n",
    "i2w = load_pickle('pickle/mov_i2w.pickle')\n",
    "\n",
    "# Reserve 0 for masking via pad_sequences\n",
    "vocab_size = len(vocab) #+ 1\n",
    "# story_maxlen = max(map(len, (x for x, _, _ in train_data + test_data)))\n",
    "query_maxlen = max(map(len, (x for _, x, _ in train_data + test_data)))\n",
    "\n",
    "print('-')\n",
    "print('Vocab size:', vocab_size, 'unique words')\n",
    "# print('Story max length:', story_maxlen, 'words')\n",
    "print('Query max length:', query_maxlen, 'words')\n",
    "print('Number of training data:', len(train_data))\n",
    "print('Number of test data:', len(test_data))\n",
    "print('-')\n",
    "print('Here\\'s what a \"data\" tuple looks like (input, query, answer):')\n",
    "print(train_data[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6101\n",
      "6101\n",
      "5909\n",
      "5909\n",
      "--------\n",
      "([], ['the movie', 'a', 'boy', 'and', 'his', 'dog,', 'when', 'was', 'it', 'released'], ['1975'])\n",
      "354\n",
      "([], ['what', 'movies', 'can', 'be', 'described', 'by', 'writer'], ['waxworks', 'the lost weekend', 'sleuth', 'the raven', 'sinister', 'the ghost writer', 'almost famous', 'the shining'])\n",
      "[] []\n"
     ]
    }
   ],
   "source": [
    "train_k_max_memsize = max(map(len, (d for d in train_k)))\n",
    "train_v_max_memsize = max(map(len, (d for d in train_v)))\n",
    "test_k_max_memsize = max(map(len, (d for d in test_k)))\n",
    "test_v_max_memsize = max(map(len, (d for d in test_v)))\n",
    "print(train_k_max_memsize)\n",
    "print(train_v_max_memsize)\n",
    "print(test_k_max_memsize)\n",
    "print(test_v_max_memsize)\n",
    "print('--------')\n",
    "# print(train_k_max_memsize, train_v_max_memsize)\n",
    "print(train_data[22281])\n",
    "len_list = []\n",
    "for i, (k,v) in enumerate(zip(train_k, train_v)):\n",
    "    if  len(k) == 0:\n",
    "        print(i)\n",
    "        print(train_data[i])\n",
    "        print(k, v)\n",
    "        break\n",
    "#     len_list.append(len(k))\n",
    "# print(sorted(len_list, reverse=True))\n",
    "\n",
    "#     print(len(k), len(v))\n",
    "#     if i > 500: break\n",
    "# #     continue\n",
    "#     if len(k) == 6041:\n",
    "#         print(i)\n",
    "#         pp.pprint(k[:10])\n",
    "#         break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "96185 9952\n",
      "95521 9878\n"
     ]
    }
   ],
   "source": [
    "train_indices, test_indices = [], []\n",
    "for i, k in enumerate(train_k):\n",
    "    if len(k) != 0:\n",
    "        train_indices.append(i)\n",
    "for i, k in enumerate(test_k):\n",
    "    if len(k) != 0:\n",
    "        test_indices.append(i)\n",
    "print(len(train_data), len(test_data))\n",
    "train_data = [train_data[i] for i in train_indices]\n",
    "train_k = [train_k[i] for i in train_indices]\n",
    "train_v = [train_v[i] for i in train_indices]\n",
    "test_data = [test_data[i] for i in test_indices]\n",
    "test_k = [test_k[i] for i in test_indices]\n",
    "test_v = [test_v[i] for i in test_indices]\n",
    "print(len(train_data), len(test_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['the movie', 'a', 'boy', 'and', 'his', 'dog,', 'when', 'was', 'it', 'released']\n"
     ]
    },
    {
     "ename": "NameError",
     "evalue": "name 'tokenize' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-5-4119e50dd8d8>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfind_ngrams\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mentities\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'the movie a boy and his dog, when was it released'\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msplit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m' '\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m100\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfind_ngrams\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mentities\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtokenize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'the movie a boy and his dog, when was it released'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m100\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfind_ngrams\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mentities\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtokenize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'w.s. van dyke was the director on which movies?'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m100\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfind_ngrams\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mentities\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mword_tokenize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'w.s. van dyke was the director on which movies?'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m100\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'tokenize' is not defined"
     ]
    }
   ],
   "source": [
    "print(find_ngrams(entities, 'the movie a boy and his dog, when was it released'.split(' '), 100))\n",
    "print(find_ngrams(entities, tokenize('the movie a boy and his dog, when was it released'), 100))\n",
    "print(find_ngrams(entities, tokenize('w.s. van dyke was the director on which movies?'), 100))\n",
    "print(find_ngrams(entities, word_tokenize('w.s. van dyke was the director on which movies?'), 100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# label_list = []\n",
    "# for _,_,a in (train_data+test_data):\n",
    "#     if a[0] not in label_list: label_list.append(a[0])\n",
    "    \n",
    "# w2i_label = dict((c, i) for i, c in enumerate(label_list))\n",
    "# i2w_label = dict((i, c) for i, c in enumerate(label_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-\n",
      "inputs: integer tensor of shape (samples, max_length)\n",
      "-\n",
      "queries: integer tensor of shape (samples, max_length)\n",
      "queries_train shape: (95521, 16)\n",
      "queries_test shape: (9878, 16)\n",
      "-\n",
      "answers: binary (1 or 0) tensor of shape (samples, vocab_size)\n",
      "answers_test shape: (9878, 74609)\n"
     ]
    }
   ],
   "source": [
    "queries_train, answers_train = vectorize(train_data,\n",
    "                                           w2i,\n",
    "                                           query_maxlen)\n",
    "queries_test, answers_test = vectorize(test_data,\n",
    "                                            w2i,\n",
    "                                            query_maxlen)\n",
    "\n",
    "print('-')\n",
    "print('inputs: integer tensor of shape (samples, max_length)')\n",
    "# print('inputs_train shape:', inputs_train.shape)\n",
    "# print('inputs_test shape:', inputs_test.shape)\n",
    "print('-')\n",
    "print('queries: integer tensor of shape (samples, max_length)')\n",
    "print('queries_train shape:', queries_train.shape)\n",
    "print('queries_test shape:', queries_test.shape)\n",
    "print('-')\n",
    "print('answers: binary (1 or 0) tensor of shape (samples, vocab_size)')\n",
    "print('answers_test shape:', answers_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vectorize_kv: 0 / 95521\n",
      "vectorize_kv: 5000 / 95521\n",
      "vectorize_kv: 10000 / 95521\n",
      "vectorize_kv: 15000 / 95521\n",
      "vectorize_kv: 20000 / 95521\n",
      "vectorize_kv: 25000 / 95521\n",
      "vectorize_kv: 30000 / 95521\n",
      "vectorize_kv: 35000 / 95521\n",
      "vectorize_kv: 40000 / 95521\n",
      "vectorize_kv: 45000 / 95521\n",
      "vectorize_kv: 50000 / 95521\n",
      "vectorize_kv: 55000 / 95521\n",
      "vectorize_kv: 60000 / 95521\n",
      "vectorize_kv: 65000 / 95521\n",
      "vectorize_kv: 70000 / 95521\n",
      "vectorize_kv: 75000 / 95521\n",
      "vectorize_kv: 80000 / 95521\n",
      "vectorize_kv: 85000 / 95521\n",
      "vectorize_kv: 90000 / 95521\n",
      "vectorize_kv: 95000 / 95521\n",
      "vectorize_kv: 0 / 95521\n",
      "vectorize_kv: 5000 / 95521\n",
      "vectorize_kv: 10000 / 95521\n",
      "vectorize_kv: 15000 / 95521\n",
      "vectorize_kv: 20000 / 95521\n",
      "vectorize_kv: 25000 / 95521\n",
      "vectorize_kv: 30000 / 95521\n",
      "vectorize_kv: 35000 / 95521\n",
      "vectorize_kv: 40000 / 95521\n",
      "vectorize_kv: 45000 / 95521\n",
      "vectorize_kv: 50000 / 95521\n",
      "vectorize_kv: 55000 / 95521\n",
      "vectorize_kv: 60000 / 95521\n",
      "vectorize_kv: 65000 / 95521\n",
      "vectorize_kv: 70000 / 95521\n",
      "vectorize_kv: 75000 / 95521\n",
      "vectorize_kv: 80000 / 95521\n",
      "vectorize_kv: 85000 / 95521\n",
      "vectorize_kv: 90000 / 95521\n",
      "vectorize_kv: 95000 / 95521\n",
      "vectorize_kv: 0 / 9878\n",
      "vectorize_kv: 5000 / 9878\n",
      "vectorize_kv: 0 / 9878\n",
      "vectorize_kv: 5000 / 9878\n"
     ]
    }
   ],
   "source": [
    "from process_data import vectorize_kv\n",
    "max_mem_len = 4\n",
    "max_mem_size = 4\n",
    "vec_train_k = vectorize_kv(train_k, max_mem_len, max_mem_size, w2i)\n",
    "vec_train_v = vectorize_kv(train_v, max_mem_len, max_mem_size, w2i)\n",
    "vec_test_k = vectorize_kv(test_k, max_mem_len, max_mem_size, w2i)\n",
    "vec_test_v = vectorize_kv(test_v, max_mem_len, max_mem_size, w2i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['which', 'movies', 'can', 'be', 'described', 'by', 'robert schwentke']\n",
      "[72580, 46780, 10202, 6723, 16929, 9914, 56132]\n",
      "[72580 46780 10202  6723 16929  9914 56132     0     0     0     0     0\n",
      "     0     0     0     0]\n",
      "['which', 'movies', 'can', 'be', 'described', 'by', 'robert schwentke', '', '', '', '', '', '', '', '', '']\n",
      "[   ['robert schwentke', '!directed_by'],\n",
      "    ['robert schwentke', '!has_tags'],\n",
      "    ['robert schwentke', '!directed_by'],\n",
      "    ['robert schwentke', '!directed_by']]\n",
      "[   [\"the time traveler's wife\"],\n",
      "    [\"the time traveler's wife\"],\n",
      "    ['flightplan'],\n",
      "    ['r.i.p.d.']]\n",
      "[[56132     1     0     0]\n",
      " [56132     4     0     0]\n",
      " [56132     1     0     0]\n",
      " [56132     1     0     0]]\n"
     ]
    },
    {
     "ename": "NameError",
     "evalue": "name 'i2w_label' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-9-3313d6abddcd>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[0mpp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_v\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mchk_id\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvec_train_k\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mchk_id\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0manswers_train\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mchk_id\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mi2w_label\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0margmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0manswers_train\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mchk_id\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'i2w_label' is not defined"
     ]
    }
   ],
   "source": [
    "chk_id = 100\n",
    "print(train_data[chk_id][1])\n",
    "print([w2i[w] for w in train_data[chk_id][1] if w in w2i])\n",
    "# print([i2w[w] for w in train_data[0][1]])\n",
    "print(queries_train[chk_id])\n",
    "print([ i2w[i] for i in queries_train[chk_id]])\n",
    "pp.pprint(train_k[chk_id])\n",
    "pp.pprint(train_v[chk_id])\n",
    "print(vec_train_k[chk_id])\n",
    "print(answers_train[chk_id], i2w_label[np.argmax(answers_train[chk_id])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 問題となる事例\n",
    "1 the movie A Boy and His Dog, when was it released?\t1975\n",
    "1 Dance, Girl, Dance, when was it released?\t1940\n",
    "Dead Man Down written_by J.H. Wyman  # word_tokenize ['Dead', 'Man', 'Down', 'written_by', 'J.H', '.', 'Wyman']\n",
    "\n",
    "load_task -> split(' 'JJ)\n",
    "kv_pairs -> word_tokenizek"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vec_test_k.shape (95521, 4, 4)\n",
      "vec_test_v.shape (95521, 4, 4)\n",
      "queries_train.shape (95521, 16)\n",
      "ans (95521, 74609)\n"
     ]
    }
   ],
   "source": [
    "from keras import backend as K\n",
    "from keras.models import Model\n",
    "from keras.layers.embeddings import Embedding\n",
    "from keras.layers import Input, Activation, Dense, Lambda, Permute, Dropout, add, multiply, dot\n",
    "\n",
    "def MemNNKV(mem_len, mem_size, query_maxlen, vocab_size, embd_size, answer_size):\n",
    "    print('mem_size:', mem_size)\n",
    "    print('q_max', query_maxlen)\n",
    "    print('embd_size', embd_size)\n",
    "    print('vocab_size', vocab_size)\n",
    "#     print('entity_size', entity_size)\n",
    "    print('-----------')\n",
    "\n",
    "    # placeholders\n",
    "    key = Input((mem_size, mem_len,), name='Key_Input')\n",
    "    val = Input((mem_size, mem_len,), name='Val_Input')\n",
    "    question = Input((query_maxlen,), name='Question_Input')\n",
    "    # print('key:', key.shape)\n",
    "\n",
    "    shared_embd_A = Embedding(input_dim=vocab_size, output_dim=embd_size)\n",
    "\n",
    "    key_encoded = shared_embd_A(key) # (None, mem_size, mem_len, embd_size)\n",
    "    # print('key_encoded', key_encoded.shape)\n",
    "    key_encoded = Lambda(lambda x: K.sum(x, axis=2)) (key_encoded) #(None, mem_size, embd_size)\n",
    "    # print('key_encoded', key_encoded.shape)\n",
    "    val_encoded = shared_embd_A(val) # (None, mem_size, embd_size)\n",
    "    val_encoded = Lambda(lambda x: K.sum(x, axis=2)) (val_encoded)\n",
    "    \n",
    "    question_encoded = shared_embd_A(question) # (None, query_max_len, embd_size)\n",
    "#     question_encoded = Lambda(lambda x: K.sum(x, axis=1)) (val_encoded)\n",
    "    question_encoded = Lambda(lambda x: K.sum(x, axis=1)) (question_encoded) #(None, embd_size)\n",
    "    # print('q_encoded', question_encoded.shape)\n",
    "    q= question_encoded\n",
    "    for h in range(2):\n",
    "        # print('---hop', h)\n",
    "        ph = dot([q, key_encoded], axes=(1, 2))  # (None, mem_size)\n",
    "        # print('ph', ph.shape)\n",
    "        ph = Activation('softmax')(ph)\n",
    "        o = multiply([ph, Permute((2, 1))(val_encoded)]) # (None, embd_size, mem_size)\n",
    "        # print('o', o.shape)\n",
    "        o = Lambda(lambda x: K.sum(x, axis=2))(o) # (None, embd_size)\n",
    "        # print('o', o.shape)\n",
    "        R = Dense(embd_size, input_shape=(embd_size,), name='R_Dense_h' + str(h+1))\n",
    "        q = R(add([q,  o])) # (None, embd_size)\n",
    "        # print('q', q.shape)\n",
    "\n",
    "#     answer = Dense(answer_size, name='last_Dense')(q) #(None, answer_size)\n",
    "    answer = Dense(vocab_size, name='last_Dense')(q) #(None, vocab_size)\n",
    "    # print('answer.shape', answer.shape)\n",
    "    preds = Activation('softmax')(answer)\n",
    "    \n",
    "    # build the final model\n",
    "    model = Model([key, val, question], preds)\n",
    "    model.compile(optimizer='adagrad', loss='categorical_crossentropy', metrics=['accuracy']) \n",
    "    return model\n",
    "\n",
    "\n",
    "print('vec_test_k.shape', vec_train_k.shape)\n",
    "print('vec_test_v.shape', vec_train_v.shape)\n",
    "print('queries_train.shape', queries_train.shape)\n",
    "print('ans', answers_train.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mem_size: 4\n",
      "q_max 16\n",
      "embd_size 500\n",
      "vocab_size 74609\n",
      "-----------\n",
      "____________________________________________________________________________________________________\n",
      "Layer (type)                     Output Shape          Param #     Connected to                     \n",
      "====================================================================================================\n",
      "Question_Input (InputLayer)      (None, 16)            0                                            \n",
      "____________________________________________________________________________________________________\n",
      "Key_Input (InputLayer)           (None, 4, 4)          0                                            \n",
      "____________________________________________________________________________________________________\n",
      "embedding_3 (Embedding)          multiple              37304500    Key_Input[0][0]                  \n",
      "                                                                   Val_Input[0][0]                  \n",
      "                                                                   Question_Input[0][0]             \n",
      "____________________________________________________________________________________________________\n",
      "Val_Input (InputLayer)           (None, 4, 4)          0                                            \n",
      "____________________________________________________________________________________________________\n",
      "lambda_13 (Lambda)               (None, 500)           0           embedding_3[2][0]                \n",
      "____________________________________________________________________________________________________\n",
      "lambda_11 (Lambda)               (None, 4, 500)        0           embedding_3[0][0]                \n",
      "____________________________________________________________________________________________________\n",
      "dot_5 (Dot)                      (None, 4)             0           lambda_13[0][0]                  \n",
      "                                                                   lambda_11[0][0]                  \n",
      "____________________________________________________________________________________________________\n",
      "lambda_12 (Lambda)               (None, 4, 500)        0           embedding_3[1][0]                \n",
      "____________________________________________________________________________________________________\n",
      "activation_7 (Activation)        (None, 4)             0           dot_5[0][0]                      \n",
      "____________________________________________________________________________________________________\n",
      "permute_5 (Permute)              (None, 500, 4)        0           lambda_12[0][0]                  \n",
      "____________________________________________________________________________________________________\n",
      "multiply_5 (Multiply)            (None, 500, 4)        0           activation_7[0][0]               \n",
      "                                                                   permute_5[0][0]                  \n",
      "____________________________________________________________________________________________________\n",
      "lambda_14 (Lambda)               (None, 500)           0           multiply_5[0][0]                 \n",
      "____________________________________________________________________________________________________\n",
      "add_5 (Add)                      (None, 500)           0           lambda_13[0][0]                  \n",
      "                                                                   lambda_14[0][0]                  \n",
      "____________________________________________________________________________________________________\n",
      "R_Dense_h1 (Dense)               (None, 500)           250500      add_5[0][0]                      \n",
      "____________________________________________________________________________________________________\n",
      "dot_6 (Dot)                      (None, 4)             0           R_Dense_h1[0][0]                 \n",
      "                                                                   lambda_11[0][0]                  \n",
      "____________________________________________________________________________________________________\n",
      "activation_8 (Activation)        (None, 4)             0           dot_6[0][0]                      \n",
      "____________________________________________________________________________________________________\n",
      "permute_6 (Permute)              (None, 500, 4)        0           lambda_12[0][0]                  \n",
      "____________________________________________________________________________________________________\n",
      "multiply_6 (Multiply)            (None, 500, 4)        0           activation_8[0][0]               \n",
      "                                                                   permute_6[0][0]                  \n",
      "____________________________________________________________________________________________________\n",
      "lambda_15 (Lambda)               (None, 500)           0           multiply_6[0][0]                 \n",
      "____________________________________________________________________________________________________\n",
      "add_6 (Add)                      (None, 500)           0           R_Dense_h1[0][0]                 \n",
      "                                                                   lambda_15[0][0]                  \n",
      "____________________________________________________________________________________________________\n",
      "R_Dense_h2 (Dense)               (None, 500)           250500      add_6[0][0]                      \n",
      "____________________________________________________________________________________________________\n",
      "last_Dense (Dense)               (None, 74609)         37379109    R_Dense_h2[0][0]                 \n",
      "____________________________________________________________________________________________________\n",
      "activation_9 (Activation)        (None, 74609)         0           last_Dense[0][0]                 \n",
      "====================================================================================================\n",
      "Total params: 75,184,609\n",
      "Trainable params: 75,184,609\n",
      "Non-trainable params: 0\n",
      "____________________________________________________________________________________________________\n",
      "None\n",
      "Train on 95521 samples, validate on 9878 samples\n",
      "Epoch 1/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 14.6887 - acc: 0.0189Epoch 00000: loss improved from inf to 14.68672, saving model to saved_models/20171018-133731_kvnn-weights-00-14.6867.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 14.6867 - acc: 0.0189 - val_loss: 8.6827 - val_acc: 0.0398\n",
      "Epoch 2/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 7.6853 - acc: 0.1155Epoch 00001: loss improved from 14.68672 to 7.68539, saving model to saved_models/20171018-133731_kvnn-weights-01-7.6854.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 7.6854 - acc: 0.1155 - val_loss: 7.2076 - val_acc: 0.1582\n",
      "Epoch 3/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 5.1807 - acc: 0.2531Epoch 00002: loss improved from 7.68539 to 5.18024, saving model to saved_models/20171018-133731_kvnn-weights-02-5.1802.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 5.1802 - acc: 0.2532 - val_loss: 6.9212 - val_acc: 0.2097\n",
      "Epoch 4/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 3.0799 - acc: 0.4278Epoch 00003: loss improved from 5.18024 to 3.07981, saving model to saved_models/20171018-133731_kvnn-weights-03-3.0798.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 3.0798 - acc: 0.4278 - val_loss: 7.0895 - val_acc: 0.2305\n",
      "Epoch 5/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 1.5879 - acc: 0.6507Epoch 00004: loss improved from 3.07981 to 1.58818, saving model to saved_models/20171018-133731_kvnn-weights-04-1.5882.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 1.5882 - acc: 0.6506 - val_loss: 7.4052 - val_acc: 0.2357\n",
      "Epoch 6/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 0.7602 - acc: 0.8336Epoch 00005: loss improved from 1.58818 to 0.76022, saving model to saved_models/20171018-133731_kvnn-weights-05-0.7602.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.7602 - acc: 0.8336 - val_loss: 7.7465 - val_acc: 0.2386\n",
      "Epoch 7/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 0.4085 - acc: 0.9121Epoch 00006: loss improved from 0.76022 to 0.40856, saving model to saved_models/20171018-133731_kvnn-weights-06-0.4086.hdf5\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "95521/95521 [==============================] - 75s - loss: 0.4086 - acc: 0.9121 - val_loss: 7.9578 - val_acc: 0.2393\n",
      "Epoch 8/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 0.2463 - acc: 0.9502Epoch 00007: loss improved from 0.40856 to 0.24626, saving model to saved_models/20171018-133731_kvnn-weights-07-0.2463.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.2463 - acc: 0.9502 - val_loss: 8.1375 - val_acc: 0.2381\n",
      "Epoch 9/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 0.1669 - acc: 0.9690Epoch 00008: loss improved from 0.24626 to 0.16696, saving model to saved_models/20171018-133731_kvnn-weights-08-0.1670.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.1670 - acc: 0.9690 - val_loss: 8.2872 - val_acc: 0.2376\n",
      "Epoch 10/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 0.1254 - acc: 0.9778Epoch 00009: loss improved from 0.16696 to 0.12532, saving model to saved_models/20171018-133731_kvnn-weights-09-0.1253.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.1253 - acc: 0.9778 - val_loss: 8.3987 - val_acc: 0.2395\n",
      "Epoch 11/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 0.1018 - acc: 0.9819Epoch 00010: loss improved from 0.12532 to 0.10188, saving model to saved_models/20171018-133731_kvnn-weights-10-0.1019.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.1019 - acc: 0.9819 - val_loss: 8.5249 - val_acc: 0.2352\n",
      "Epoch 12/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 0.0874 - acc: 0.9850Epoch 00011: loss improved from 0.10188 to 0.08742, saving model to saved_models/20171018-133731_kvnn-weights-11-0.0874.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0874 - acc: 0.9850 - val_loss: 8.5823 - val_acc: 0.2377\n",
      "Epoch 13/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 0.0787 - acc: 0.9861Epoch 00012: loss improved from 0.08742 to 0.07880, saving model to saved_models/20171018-133731_kvnn-weights-12-0.0788.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0788 - acc: 0.9860 - val_loss: 8.6457 - val_acc: 0.2390\n",
      "Epoch 14/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 0.0712 - acc: 0.9876Epoch 00013: loss improved from 0.07880 to 0.07119, saving model to saved_models/20171018-133731_kvnn-weights-13-0.0712.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0712 - acc: 0.9876 - val_loss: 8.7096 - val_acc: 0.2350\n",
      "Epoch 15/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 0.0660 - acc: 0.9886Epoch 00014: loss improved from 0.07119 to 0.06603, saving model to saved_models/20171018-133731_kvnn-weights-14-0.0660.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0660 - acc: 0.9886 - val_loss: 8.7642 - val_acc: 0.2373\n",
      "Epoch 16/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 0.0631 - acc: 0.9887Epoch 00015: loss improved from 0.06603 to 0.06308, saving model to saved_models/20171018-133731_kvnn-weights-15-0.0631.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0631 - acc: 0.9887 - val_loss: 8.7952 - val_acc: 0.2371\n",
      "Epoch 17/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 0.0608 - acc: 0.9892Epoch 00016: loss improved from 0.06308 to 0.06080, saving model to saved_models/20171018-133731_kvnn-weights-16-0.0608.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0608 - acc: 0.9892 - val_loss: 8.8302 - val_acc: 0.2377\n",
      "Epoch 18/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 0.0577 - acc: 0.9895Epoch 00017: loss improved from 0.06080 to 0.05766, saving model to saved_models/20171018-133731_kvnn-weights-17-0.0577.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0577 - acc: 0.9895 - val_loss: 8.8729 - val_acc: 0.2363\n",
      "Epoch 19/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 0.0560 - acc: 0.9895Epoch 00018: loss improved from 0.05766 to 0.05602, saving model to saved_models/20171018-133731_kvnn-weights-18-0.0560.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0560 - acc: 0.9895 - val_loss: 8.8846 - val_acc: 0.2353\n",
      "Epoch 20/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 0.0538 - acc: 0.9903Epoch 00019: loss improved from 0.05602 to 0.05376, saving model to saved_models/20171018-133731_kvnn-weights-19-0.0538.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0538 - acc: 0.9903 - val_loss: 8.9027 - val_acc: 0.2380\n",
      "Epoch 21/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 0.0529 - acc: 0.9902Epoch 00020: loss improved from 0.05376 to 0.05292, saving model to saved_models/20171018-133731_kvnn-weights-20-0.0529.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0529 - acc: 0.9902 - val_loss: 8.9532 - val_acc: 0.2333\n",
      "Epoch 22/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 0.0506 - acc: 0.9908Epoch 00021: loss improved from 0.05292 to 0.05060, saving model to saved_models/20171018-133731_kvnn-weights-21-0.0506.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0506 - acc: 0.9908 - val_loss: 8.9654 - val_acc: 0.2376\n",
      "Epoch 23/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 0.0499 - acc: 0.9909Epoch 00022: loss improved from 0.05060 to 0.04991, saving model to saved_models/20171018-133731_kvnn-weights-22-0.0499.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0499 - acc: 0.9909 - val_loss: 8.9958 - val_acc: 0.2371\n",
      "Epoch 24/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 0.0488 - acc: 0.9912Epoch 00023: loss improved from 0.04991 to 0.04876, saving model to saved_models/20171018-133731_kvnn-weights-23-0.0488.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0488 - acc: 0.9912 - val_loss: 8.9963 - val_acc: 0.2376\n",
      "Epoch 25/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 0.0478 - acc: 0.9913Epoch 00024: loss improved from 0.04876 to 0.04791, saving model to saved_models/20171018-133731_kvnn-weights-24-0.0479.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0479 - acc: 0.9913 - val_loss: 9.0209 - val_acc: 0.2373\n",
      "Epoch 26/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 0.0463 - acc: 0.9915Epoch 00025: loss improved from 0.04791 to 0.04624, saving model to saved_models/20171018-133731_kvnn-weights-25-0.0462.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0462 - acc: 0.9915 - val_loss: 9.0535 - val_acc: 0.2363\n",
      "Epoch 27/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 0.0460 - acc: 0.9918Epoch 00026: loss improved from 0.04624 to 0.04599, saving model to saved_models/20171018-133731_kvnn-weights-26-0.0460.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0460 - acc: 0.9918 - val_loss: 9.0505 - val_acc: 0.2385\n",
      "Epoch 28/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 0.0451 - acc: 0.9918Epoch 00027: loss improved from 0.04599 to 0.04513, saving model to saved_models/20171018-133731_kvnn-weights-27-0.0451.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0451 - acc: 0.9918 - val_loss: 9.0719 - val_acc: 0.2358\n",
      "Epoch 29/30\n",
      "95424/95521 [============================>.] - ETA: 0s - loss: 0.0442 - acc: 0.9921Epoch 00028: loss improved from 0.04513 to 0.04421, saving model to saved_models/20171018-133731_kvnn-weights-28-0.0442.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0442 - acc: 0.9921 - val_loss: 9.0953 - val_acc: 0.2370\n",
      "Epoch 30/30\n",
      "95488/95521 [============================>.] - ETA: 0s - loss: 0.0437 - acc: 0.9921Epoch 00029: loss improved from 0.04421 to 0.04375, saving model to saved_models/20171018-133731_kvnn-weights-29-0.0437.hdf5\n",
      "95521/95521 [==============================] - 75s - loss: 0.0437 - acc: 0.9921 - val_loss: 9.1045 - val_acc: 0.2368\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7fe06c017eb8>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embd_size = 500\n",
    "memnn_kv = MemNNKV(max_mem_len, max_mem_size, query_maxlen, vocab_size, embd_size, None)\n",
    "print(memnn_kv.summary())\n",
    "now = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')\n",
    "model_path = 'saved_models/' + now + '_kvnn-weights-{epoch:02d}-{loss:.4f}.hdf5'\n",
    "checkpoint = ModelCheckpoint(model_path, monitor='loss', verbose=1, save_best_only=True, mode='min')\n",
    "callbacks_list = [checkpoint]\n",
    "memnn_kv.fit([vec_train_k, vec_train_v, queries_train], answers_train,\n",
    "          batch_size=64,\n",
    "          epochs=30,\n",
    "          callbacks=callbacks_list,\n",
    "          validation_data=([vec_test_k, vec_test_v, queries_test], answers_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "memnn_kv.save('model_memnn_kv.h5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "model = load_model('demo_model_memnn_kv.h5') # heavy to run\n",
    "# score = model.evaluate([vec_test_k, vec_test_v, queries_test], answers_test, verbose=1)\n",
    "# score = model.evaluate([vec_train_k, vec_train_v, queries_train], answers_train, verbose=1)\n",
    "# print('Test loss:', score[0])\n",
    "# print('Test accuracy:', score[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# pred = model.predict([vec_train_k, vec_train_v, queries_train], verbose=1) \n",
    "pred = load_pickle('pred.pickle')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "for i, (p, a) in enumerate(zip(pred, answers_train[:len(pred)])):\n",
    "    pw = i2w[np.argmax(p)]\n",
    "    pa = i2w[np.argmax(a)]\n",
    "    if pw == pa:\n",
    "        print(i, ' '.join(train_data[i][1]))\n",
    "        print(pw)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print('q.shape', queries_test.shape)\n",
    "ret_predict = model.predict([vec_test_k, vec_test_v, queries_test], verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print('q:',test_data[0][1])\n",
    "print('predict:',i2w[np.argmax(ret_predict[0])])\n",
    "print('label:',i2w[np.argmax(answers_test[0])])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interactive\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# stopwords = load_pickle('mov_stopwords.pickle')\n",
    "\n",
    "# tokenize a question\n",
    "q = 'who directed blade runner'\n",
    "q_tokens = word_tokenize(q)\n",
    "q_tokens = find_ngrams(vocab, q_tokens, 100000)\n",
    "print('q_tokens:', q_tokens)\n",
    "\n",
    "# vectorize a question\n",
    "q_pad_len = max(0, query_maxlen - len(q_tokens))\n",
    "vec_q = [w2i[w] for w in q_tokens] + [0] * q_pad_len\n",
    "vec_q = np.array(vec_q)\n",
    "vec_q = np.reshape(vec_q, (1, len(vec_q)))\n",
    "print('vec_q:', vec_q)\n",
    "\n",
    "# get related kv\n",
    "k_list, v_list = [], []\n",
    "for w in q_tokens:\n",
    "    if w not in stopwords:\n",
    "        for kv_ind, (k, v) in enumerate(kv_pairs):\n",
    "            if w in (k+v):\n",
    "                k_list += k\n",
    "                v_list += v\n",
    "    else:\n",
    "        print(w, 'in stopwords')\n",
    "\n",
    "def _vec_kv(data, w2i, mem_maxlen):\n",
    "    vec = [w2i[e] for e in data if e in w2i]\n",
    "    vec += [0] * max(0, mem_maxlen - len(vec))\n",
    "    vec = vec[:mem_maxlen]\n",
    "    vec = np.array(vec)\n",
    "#     vec = np.expand_dims(vec, axis=0)\n",
    "    print('---', vec.shape)\n",
    "    vec = np.reshape(vec, (1, 100))\n",
    "    \n",
    "    return vec\n",
    "# print(k_list)\n",
    "# vectroize kv\n",
    "vec_k, vec_v = None, None\n",
    "vec_k = _vec_kv(k_list, w2i, mem_maxlen)\n",
    "vec_v = _vec_kv(v_list, w2i, mem_maxlen)\n",
    "# vec_k = [w2i[e] for e in k_list if e in w2i]\n",
    "# vec_k += [0] * max(0, mem_maxlen - len(vec_k))\n",
    "# vec_k = vec_k[:mem_maxlen]\n",
    "# vec_k = np.array(vec_k)\n",
    "# vec_v = [w2i[e] for e in v_list if e in w2i]\n",
    "# vec_v += [0] * max(0, mem_maxlen - len(vec_v))\n",
    "# vec_v = vec_v[:mem_maxlen]\n",
    "# vec_v = np.array(vec_v)\n",
    "print(vec_k.shape)\n",
    "\n",
    "int_predict = model.predict([vec_k, vec_v, vec_q], batch_size=1, verbose=1)     \n",
    "print('q:',q)\n",
    "print('predict:',i2w[np.argmax(int_predict[0])])\n",
    "# print('label:',i2w[np.argmax(answers_test[0])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## misc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
